
import asyncio
import os
import sys
import json
from typing import Optional, List, Tuple
from contextlib import AsyncExitStack

from openai import OpenAI
from dotenv import load_dotenv

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

# 加载 .env 文件，确保 API Key 受到保护
load_dotenv()

SUPPORTED_EXTS = (".pdf", ".txt", ".docx")

def scan_documents(base_path: str, extensions: Tuple[str, ...] = SUPPORTED_EXTS) -> list[str]:
    """
    递归扫描 base_path 下所有指定扩展名的文档，并返回绝对路径列表（排序稳定）。
    """
    file_paths: List[str] = []
    for root, _, files in os.walk(base_path):
        for file in files:
            if file.lower().endswith(extensions):
                file_paths.append(os.path.abspath(os.path.join(root, file)))
    # 稳定排序，保证展示一致
    file_paths.sort()
    return file_paths

def render_docs_table(docs: List[str]) -> str:
    """
    将扫描到的文档渲染为带序号的清单，供展示和系统提示。
    """
    if not docs:
        return "（未扫描到任何文档）"
    lines = []
    for i, p in enumerate(docs, 1):
        lines.append(f"{i}. {p}")
    return "\n".join(lines)

class MCPClient:
    def __init__(self):
        """初始化 MCP 客户端"""
        self.exit_stack = AsyncExitStack()
        self.openai_api_key = os.getenv("OPENAI_API_KEY")  # 读取 OpenAI API Key
        self.base_url = os.getenv("BASE_URL")  # 读取 BASE URL
        self.model = os.getenv("MODEL") or "deepseek-chat"  # 读取 model
        if not self.openai_api_key:
            raise ValueError("❌ 未找到 OpenAI API Key，请在 .env 文件中设置 OPENAI_API_KEY")
        self.client = OpenAI(api_key=self.openai_api_key, base_url=self.base_url)  # 创建OpenAI client
        self.session: Optional[ClientSession] = None
        # 扫描配置：支持通过环境变量覆盖
        self.docs_dir = os.getenv("DOCS_DIR", "./documents")
        self.docs_exts = tuple(e.strip().lower() for e in os.getenv("DOCS_EXTS", ".pdf,.txt,.docx").split(","))
        # 结果缓存
        self.scanned_docs: List[str] = []

    async def connect_to_server(self, server_script_path: str):
        """连接到 MCP 服务器并列出可用工具"""
        is_python = server_script_path.endswith('.py')
        is_js = server_script_path.endswith('.js')
        if not (is_python or is_js):
            raise ValueError("服务器脚本必须是 .py 或 .js 文件")

        command = "python" if is_python else "node"
        server_params = StdioServerParameters(
            command=command,
            args=[server_script_path],
            env=None
        )

        # 启动 MCP 服务器并建立通信
        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.stdio, self.write = stdio_transport
        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))

        await self.session.initialize()

        # 列出 MCP 服务器上的工具
        response = await self.session.list_tools()
        tools = response.tools
        print("\n已连接到服务器，支持以下工具:", [tool.name for tool in tools])

    def rescan(self):
        """重新扫描文档目录。"""
        print(f"📂 正在扫描目录: {os.path.abspath(self.docs_dir)}")
        self.scanned_docs = scan_documents(self.docs_dir, self.docs_exts)
        if self.scanned_docs:
            print("✅ 已找到以下文档：")
            for i, doc in enumerate(self.scanned_docs, 1):
                print(f"{i}. {doc}")
        else:
            print("⚠️ 未找到文档（可修改 DOCS_DIR 或 DOCS_EXTS 后重新扫描）。")

    def build_system_prompt(self) -> str:
        """构造系统提示，把可用文档清单注入，指示模型如何自动选择文件路径。"""
        docs_table = render_docs_table(self.scanned_docs)
        system = f"""你是一个多功能助手，可以查询天气和回答文档相关问题。

工具使用规则：

天气查询：
1. 当用户询问任何与天气相关的问题时，**必须**调用 query_weather 工具。
2. 工具支持中文城市名称，如 "北京"、"西安"、"上海" 等，如果用户输入中文，你需要先将中文转化成英文名称。
3. 如果用户没有明确指定城市，请询问具体城市。

文档处理：
1. 当用户询问关于文档内容的问题时，**必须**使用 read_and_query_document 工具；如果用户要求摘要，使用 get_document_summary 工具。
2. 可供选择的文档文件（从本地自动扫描目录得到）：
{docs_table}
3. 当用户未提供具体路径时：
   - 根据用户描述的关键词（标题、目录名、文件名关键字），从上述清单中**自动选择最匹配的一条绝对路径**作为 file_path 参数；
   - 若匹配度相近或不明确，先向用户确认候选（最多 5 项），再调用工具；
   - 如果清单为空，请先提示用户把文件放到扫描目录或更改扫描路径。

其它：
- 所有与文档有关的调用都应优先使用上方清单中的**绝对路径**；
- 若用户使用“第 N 个”、“编号 N”这类指代，请将其解析为清单中的相应绝对路径；
- 回答后尽量简短，必要时再询问澄清。"""
        return system

    async def process_query(self, messages: list) -> str:
        """
        使用大模型处理查询并调用可用的 MCP 工具 (Function Calling)，支持多个工具调用
        """
        max_iterations = 5  # 防止无限循环
        iteration = 0
        while iteration < max_iterations:
            iteration += 1

            # 获取可用工具列表
            response = await self.session.list_tools()
            available_tools = [{
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "input_schema": tool.inputSchema
                }
            } for tool in response.tools]

            # 调用大模型
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                tools=available_tools
            )
            choice = response.choices[0]
            content = choice.message

            if not content.tool_calls:
                return content.content or ""

            # 将模型的消息添加到历史中
            messages.append(content.model_dump())

            # 并行执行所有工具调用
            tool_results = []
            for tool_call in content.tool_calls:
                tool_name = tool_call.function.name
                try:
                    tool_args = json.loads(tool_call.function.arguments or "{}")
                except Exception:
                    tool_args = {}
                print(f"[调用工具 {tool_name} 参数 {tool_args}]")
                try:
                    result = await self.session.call_tool(tool_name, tool_args)
                    tool_results.append({
                        "tool_call_id": tool_call.id,
                        "result": result.content[0].text if result.content else ""
                    })
                except Exception as e:
                    tool_results.append({
                        "tool_call_id": tool_call.id,
                        "result": f"工具调用失败: {str(e)}"
                    })

            for result in tool_results:
                messages.append({
                    "role": "tool",
                    "content": result["result"],
                    "tool_call_id": result["tool_call_id"],
                })

            if iteration == max_iterations:
                final_response = self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                )
                return final_response.choices[0].message.content or ""
        return "达到最大迭代次数，请简化您的请求。"

    async def chat_loop(self):
        """运行交互式聊天循环"""
        print("\n🤖 MCP 客户端已启动！输入 'quit' 退出")
        print("💡 命令：'docs' 重新列出文档，'rescan' 重新扫描，'quit' 退出\n")

        # 初次扫描
        self.rescan()

        # 构造带有文档清单的系统提示
        conversation_history = [
            {"role": "system", "content": self.build_system_prompt()}
        ]

        while True:
            try:
                query = input("\n你: ").strip()
                if not query:
                    continue
                if query.lower() == 'quit':
                    break
                if query.lower() == 'docs':
                    print("\n📄 当前可用文档：")
                    print(render_docs_table(self.scanned_docs))
                    continue
                if query.lower() == 'rescan':
                    self.rescan()
                    # 更新系统提示中的文档清单
                    conversation_history[0]["content"] = self.build_system_prompt()
                    continue

                # 将“第 N 个”指代解析为路径注入用户问题的上下文提示
                # 例：用户输入 “总结第2个文档” -> 在消息中追加一个轻量指示
                idx_hint = None
                import re
                m = re.search(r"(第\s*(\d+)\s*个)", query)
                if m:
                    try:
                        idx = int(m.group(2))
                        if 1 <= idx <= len(self.scanned_docs):
                            idx_hint = self.scanned_docs[idx - 1]
                    except Exception:
                        pass

                if idx_hint:
                    query = f"{query}\n（如果需要文件路径，请使用：{idx_hint}）"

                conversation_history.append({"role": "user", "content": query})

                response = await self.process_query(conversation_history)  # 发送用户输入到 OpenAI API
                conversation_history.append({"role": "assistant", "content": response})
                print(f"\n🤖 OpenAI: {response}")

                # 控制历史长度避免 token 过多（保留系统 + 最近 9 轮）
                if len(conversation_history) > 1 + 18:
                    # 留下 system 和最近 18 条
                    conversation_history = [conversation_history[0]] + conversation_history[-18:]

            except Exception as e:
                print(f"\n⚠️ 发生错误: {str(e)}")

    async def cleanup(self):
        """清理资源"""
        await self.exit_stack.aclose()


async def main():
    if len(sys.argv) < 2:
        print("Usage: python client_scan.py <path_to_server_script>")
        sys.exit(1)

    client = MCPClient()
    try:
        await client.connect_to_server(sys.argv[1])
        await client.chat_loop()
    finally:
        await client.cleanup()


if __name__ == "__main__":
    asyncio.run(main())
